
import numpy as np
from scipy.interpolate import CubicSpline


points_v3 = np.array([[0.0, 0, 0.0, 0, 0, 0],
                   [0.0, 0, 3.3*0.5, 0, 0, 0],
                   [0.0, 0, 3.3, 0, 0, 0],
                   [2.0, 0, 3.3, 0, 0, 0],
                   [4.0, 0, 3.3, 0, 0, 0],
                   [4.0, 0, 3.3*0.5, 0, 0, 0],
                   [4.0, 0, 0.0, 0, 0, 0]],dtype= np.float16)

points_v1 = np.array([[0.0, 0, 0.0, 0, 0, 0],
                   [0.0, 0, 3.3*0.5, 0, 0, 0],
                   [0.0, 0, 3.3, 0, 0, 0],
                   [2.0, 0, 3.3, 0, 0, 0],
                   [4.0, 0, 3.3, 0, 0, 0],
                   [4.0, 0, 3.3, 0, 0, -0.25*np.pi],
                   [4.0, 0, 3.3, 0, 0, -0.5*np.pi],
                   [4.0, 0, 3.3, 0, 0, -0.75*np.pi],
                   [4.0, 0, 3.3, 0, 0, -np.pi],
                   [4.0, 0, 3.3, 0, 0, 0.75*np.pi],
                   [4.0, 0, 3.3, 0, 0, 0.5*np.pi],
                   [4.0, 0, 3.3, 0, 0, 0.25*np.pi],
                   [4.0, 0, 3.3, 0, 0, 0],
                   [4.0, 0, 3.3*0.5, 0, 0, 0],
                   [4.0, 0, 0.0, 0, 0, 0]],dtype= np.float16)

def get_sparse_indices(sparse_traj, number=40):
    # 分解为x、y、z三个数组
    t = np.linspace(0, 1, len(sparse_traj))  # 原始点的参数
    t_fine = np.linspace(0, 1, number)  # 细化的参数

    # 找到每个稀疏点的对应索引
    indices = [np.argmin(np.abs(t_fine - ti)) for ti in t]
    return indices

def get_traj(sparse_traj, number = 40):
    # 给定的四个点

    # 分解为x、y、z三个数组
    x = sparse_traj[:, 0]
    y = sparse_traj[:, 1]
    z = sparse_traj[:, 2]

    # 生成参数t
    t = np.linspace(0, 1, len(sparse_traj))

    # 对每个坐标使用样条插值生成参数化曲线
    cs_x = CubicSpline(t, x)
    cs_y = CubicSpline(t, y)
    cs_z = CubicSpline(t, z)

    # 生成用于绘制的细化参数
    t_fine = np.linspace(0, 1, number)

    # 计算细化参数下的曲线
    x_fine = cs_x(t_fine)
    y_fine = cs_y(t_fine)
    z_fine = cs_z(t_fine)

    traj = np.array([x_fine,y_fine,z_fine]).T
    return traj
    

points_v1_for_interpolation = np.array([[0.0, 0, 0.0, 0, 0, 0],
                   [0.0, 0, 3.3*0.5, 0, 0, 0],
                   [0.0, 0, 3.3, 0, 0, 0],
                   [2.0, 0, 3.3, 0, 0, 0],
                   [4.0, 0, 3.3, 0, 0, 0],
                   [4.0, 0, 3.3, 0, 0, -0.5*np.pi],
                   [4.0, 0, 3.3, 0, 0, -np.pi],
                   [4.0, 0, 3.3, 0, 0, -1.5*np.pi],
                   [4.0, 0, 3.3, 0, 0, -2*np.pi],
                   [4.0, 0, 3.3*0.5, 0, 0, -2*np.pi],
                   [4.0, 0, 0.0, 0, 0, -2*np.pi]],dtype= np.float16)


def get_traj_pose(sparse_traj, number=40):
    """
    生成细化的轨迹，包括位置 (x, y, z) 和姿态 (roll, pitch, yaw)，
    对每个坐标使用样条插值生成参数化曲线，且解决周期性角度插值问题。

    参数：
    sparse_traj (np.ndarray): 输入的稀疏轨迹 (N, 6)，包括 x, y, z, roll, pitch, yaw
    number (int): 生成细化轨迹的点数
    
    返回：
    np.ndarray: 细化轨迹 (number, 6)，包括 x, y, z, roll, pitch, yaw
    """
    # 分解为 x, y, z 和 r, p, y（分别代表 roll, pitch, yaw）
    x = sparse_traj[:, 0]
    y = sparse_traj[:, 1]
    z = sparse_traj[:, 2]
    roll = sparse_traj[:, 3]
    pitch = sparse_traj[:, 4]
    yaw = sparse_traj[:, 5]

    # 生成参数 t
    t = np.linspace(0, 1, len(sparse_traj))

    # 对每个坐标和角度使用样条插值生成参数化曲线
    cs_x = CubicSpline(t, x)
    cs_y = CubicSpline(t, y)
    cs_z = CubicSpline(t, z)
    cs_roll = CubicSpline(t, roll)
    cs_pitch = CubicSpline(t, pitch)
    cs_yaw = CubicSpline(t, yaw)

    # 生成细化的参数 t_fine
    t_fine = np.linspace(0, 1, number)

    # 计算细化参数下的轨迹
    x_fine = cs_x(t_fine)
    y_fine = cs_y(t_fine)
    z_fine = cs_z(t_fine)
    roll_fine = cs_roll(t_fine)
    pitch_fine = cs_pitch(t_fine)
    yaw_fine = cs_yaw(t_fine)



    traj = np.vstack([x_fine, y_fine, z_fine, roll_fine, pitch_fine, yaw_fine]).T

    return traj

def find_closest_indices(dense_traj, sparse_traj):
        """
        找出 dense_traj 中与 sparse_traj 对应的索引。
        
        参数：
        dense_traj (np.ndarray): 细化后的轨迹数组 (N, 6)
        sparse_traj (np.ndarray): 原始轨迹数组 (M, 6)
        
        返回：
        np.ndarray: sparse_traj 中每个点在 dense_traj 中对应的索引
        """
        indices = []
        for point in sparse_traj:
            # 计算 dense_traj 中每个点与 current point 的欧几里得距离
            distances = np.linalg.norm(dense_traj[:, :3] - point[:3], axis=1)  # 只计算位置 (x, y, z) 的距离
            # 找到距离最小的点的索引
            index = np.argmin(distances)
            indices.append(index)
        
        return np.array(indices)
def get_sparse_pose_indices(sparse_traj, number=40):
    
    dense_traj = get_traj_pose(sparse_traj, number)
    # 假设 dense_traj 和 points_v1 已经定义好
    indices = find_closest_indices(dense_traj, sparse_traj)
    return indices

def generate_equally_spaced_points(min_point, max_point, num_points_per_dim):
    """
    在给定三维矩形区域内均匀等间隔生成点。
    
    参数：
    min_point (tuple): 矩形区域的最小点 (x_min, y_min, z_min)
    max_point (tuple): 矩形区域的最大点 (x_max, y_max, z_max)
    num_points_per_dim (int): 每个维度上生成的点数
    
    返回：
    np.ndarray: 生成的三维点数组
    """
    # 在每个维度上生成等间隔的点
    x = np.linspace(min_point[0], max_point[0], num_points_per_dim)
    y = np.linspace(min_point[1], max_point[1], num_points_per_dim)
    z = np.linspace(min_point[2], max_point[2], num_points_per_dim)
    
    # 使用 meshgrid 生成所有的组合
    X, Y, Z = np.meshgrid(x, y, z)
    
    # 将 X, Y, Z 合并为一个二维数组，其中每一行是一个点的坐标
    points = np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T
    
    return points

# 定义矩形区域的最小值和最大值
min_point = (-1.75, -2.4, 0)
max_point = (4.75, 2.4, 3.3)
num_points_per_dim = 5  
points_v2 = generate_equally_spaced_points(min_point, max_point, num_points_per_dim)
orients = np.zeros_like(points_v2)
points_v2 = np.hstack([points_v2, orients])


